import numpy as np
from assignment2.python import helper


def gen(dot1, dot):
    return np.array([
        dot1[0] * dot[0],
        dot1[0] * dot[1],
        dot1[0],
        dot1[1] * dot[0],
        dot1[1] * dot[1],
        dot1[1],
        dot[0],
        dot[1],
        1], dtype=np.float32)


# 编写八点算法，pts是N*2矩阵，两张图中的特征点坐标
def eight_point(Pts1, Pts2, M):
    # get the number of pairs
    num = Pts1.shape[0]
    height, width = M
    T = np.array([[1 / width, 0, 0], [0, 1 / height, 0], [0, 0, 1]], dtype=np.float32)
    add = np.ones((1, Pts1.shape[0]), dtype=np.float32, order='C')
    pts1 = np.insert(Pts1, 2, values=add, axis=1)
    pts2 = np.insert(Pts2, 2, values=add, axis=1)
    # print(pts1)

    # transform the dots to range(0,1)
    pts1 = (T @ pts1.T).T
    pts2 = (T @ pts2.T).T
    # print(pts1)

    # create m*9 matrix
    A = gen(pts1[0], pts2[0])
    for i in range(1, num):
        # print(A)
        A = np.append(A, values=gen(pts1[i], pts2[i]), axis=0)
        # print(A)
    A = np.reshape(A, (num, 9))
    # print(A)

    # SVD process and get initial F
    U, S, V = np.linalg.svd(A)
    # f = np.reshape(V.T[-1], (3, 3))
    f = V[:, -1].reshape(3, 3)

    # transform rank of f to 2
    u, s, v = np.linalg.svd(f)
    s1 = np.zeros((3, 3), dtype=np.float32)
    # print(s) 3*1 matrix
    s1[0][0] = s[0]
    s1[1][1] = s[1]
    # print(s1) 3*3 matrix
    F = u @ s1 @ v

    # refining with provided function
    F = helper.refineF(F, Pts1, Pts2)

    # back
    F = T.T @ F @ T
    #print(F)

    return F


def FindFundamentalMatrix(Pts1, Pts2):
    # Normalize the points&transformation m   (实验手册中的M加入后效果太差，因此重新完整写了一个)
    pts1, pts2 = Pts1.T, Pts2.T  # (2,N)
    n = pts1.shape[1]
    m1, m2 = np.mean(pts1, axis=1)[:, np.newaxis], np.mean(pts2, axis=1)[:, np.newaxis]  # mean
    d1, d2 = np.mean(np.sum((pts1 - m1) ** 2, axis=0) ** 0.5), np.mean(
        np.sum((pts2 - m2) ** 2, axis=0) ** 0.5)  # mean distance
    s1, s2 = 2 ** 0.5 / d1, 2 ** 0.5 / d2  # scales

    T1 = np.array([[s1, 0, -m1[0, 0] * s1],
                   [0, s1, -m1[1, 0] * s1],
                   [0, 0, 1]])
    T2 = np.array([[s2, 0, -m2[0, 0] * s2],
                   [0, s2, -m2[1, 0] * s2],
                   [0, 0, 1]])
    pts1, pts2 = np.vstack([pts1, np.array([1] * n)]), np.vstack([pts2, np.array([1] * n)])  # (3,N) matrix
    pts1_n, pts2_n = (T1 @ pts1).T, (T2 @ pts2).T

    # the matrix A
    A = np.zeros((n, 9))
    for i in range(n):
        p1, p2 = pts1_n[i], pts2_n[i]
        x1, y1, x2, y2 = p1[0], p1[1], p2[0], p2[1]
        A[i] = np.array([x2 * x1, x2 * y1, x2, y2 * x1, y2 * y1, y2, x1, y1, 1.0])

    # Find the fundamental matrix
    u, s, v = np.linalg.svd(A)

    # take the vector wrt the smallest singular value
    F = v[-1].reshape(3, 3)

    # change rank(F) to 2
    u, s, v = np.linalg.svd(F)
    F = u @ np.diag([*s[:2], 0]) @ v

    # F = helper.refineF(F, Pts1, Pts2)
    # weird：不要这个优化，反而效果更好
    F = T2.T @ F @ T1
    return F
